{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": " Unfolded_WMMSE_versus_WMMSE.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X6NBSyvW9e5Q",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "# Unfolded WMMSE versus WMMSE\n",
        "The WMMSE algorithm is implemented as described in the paper \"An Iteratively Weighted MMSE Approach to Distributed Sum-Utility Maximization for \n",
        "a MIMO Interfering Broadcast Channel\", by Q. Shi et al. \n",
        "\n",
        "Note:\n",
        "*   Functions and variables with \"_nn\" (i.e. neural network) in the name refer to the unfolded WMMSE algorithm\n",
        "*   In the code we use the term precoder to indicate the beamformer \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xF5eK6Qt-J7g",
        "colab_type": "text"
      },
      "source": [
        "### Import libraries and set variables"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cRHBPbJT9_dN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Import libraries\n",
        "import tensorflow as tf # tensorflow_version 1.x is needed\n",
        "import numpy as np\n",
        "import copy\n",
        "from copy import deepcopy\n",
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "# Set variables \n",
        "nr_of_users = 4\n",
        "nr_of_BS_antennas = 4\n",
        "scheduled_users = [0,1,2,3] # array of scheduled users. Note that we schedule all the users.\n",
        "epsilon = 0.0001 # used to end the iterations of the WMMSE algorithm in Shi et al. when the number of iterations is not fixed (note that the stopping criterion has precendence over the fixed number of iterations)\n",
        "power_tolerance = 0.0001 # used to end the bisection search in the WMMSE algorithm in Shi et al.\n",
        "total_power = 100 # power constraint in the weighted sum rate maximization problem - eq. (4) in our paper\n",
        "noise_power = 1\n",
        "nr_of_batches_training = 10000 # used for training\n",
        "nr_of_batches_test = 1000 # used for testing\n",
        "nr_of_samples_per_batch = 100 \n",
        "nr_of_iterations = 1 # for WMMSE algorithm in Shi et al. \n",
        "nr_of_iterations_nn = 1 # for the unfolded WMMSE in our paper\n",
        "\n",
        "# User weights in the weighted sum rate (denoted by alpha in eq. (4) in our paper)\n",
        "user_weights = np.reshape(np.ones(nr_of_users*nr_of_samples_per_batch),(nr_of_samples_per_batch,nr_of_users,1))\n",
        "user_weights_for_regular_WMMSE = np.ones(nr_of_users)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2w-i4Wq5-FAV",
        "colab_type": "text"
      },
      "source": [
        "### Define utility functions\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vyYStVlV9f-w",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Compute power for bisection search in the optimization of the transmitter precoder \n",
        "# - eq. (18) in the paper by Shi et al.\n",
        "def compute_P(Phi_diag_elements, Sigma_diag_elements, mu):\n",
        "  nr_of_BS_antennas = Phi_diag_elements.size\n",
        "  mu_array = mu*np.ones(Phi_diag_elements.size)\n",
        "  result = np.divide(Phi_diag_elements,(Sigma_diag_elements + mu_array)**2)\n",
        "  result = np.sum(result)\n",
        "  return result\n",
        "\n",
        "\n",
        "def compute_norm_of_complex_array(x):\n",
        "  result = np.sqrt(np.sum((np.absolute(x))**2))\n",
        "  return result\n",
        "\n",
        "\n",
        "def compute_sinr(channel, precoder, noise_power, user_id, selected_users):\n",
        "    nr_of_users = np.size(channel,0)\n",
        "    numerator = (np.absolute(np.matmul(np.conj(channel[user_id,:]),precoder[user_id,:])))**2\n",
        "\n",
        "    inter_user_interference = 0\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index != user_id and user_index in selected_users:\n",
        "        inter_user_interference = inter_user_interference + (np.absolute(np.matmul(np.conj(channel[user_id,:]),precoder[user_index,:])))**2\n",
        "    denominator = noise_power + inter_user_interference\n",
        "\n",
        "    result = numerator/denominator\n",
        "    return result\n",
        "\n",
        "\n",
        "def compute_user_weights(nr_of_users, selected_users):\n",
        "  result = np.ones(nr_of_users)\n",
        "  for user_index in range(nr_of_users):\n",
        "    if not (user_index in selected_users):\n",
        "      result[user_index] = 0\n",
        "  return result\n",
        "\n",
        "\n",
        "def compute_weighted_sum_rate(user_weights, channel, precoder, noise_power, selected_users):\n",
        "   result = 0\n",
        "   nr_of_users = np.size(channel,0)\n",
        "   \n",
        "   for user_index in range(nr_of_users):\n",
        "     if user_index in selected_users:\n",
        "       user_sinr = compute_sinr(channel, precoder, noise_power, user_index, selected_users)\n",
        "       result = result + user_weights[user_index]*np.log2(1 + user_sinr)\n",
        "    \n",
        "   return result\n",
        "\n",
        "\n",
        "def compute_sinr_nn(channel, precoder, noise_power, user_id, nr_of_users):\n",
        "\n",
        "    numerator = tf.reduce_sum((tf.matmul(tf.transpose(channel[user_id]),precoder[user_id]))**2)\n",
        "    inter_user_interference = 0\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index != user_id:\n",
        "        inter_user_interference = inter_user_interference +  tf.reduce_sum((tf.matmul(tf.transpose(channel[user_id]),precoder[user_index]))**2)\n",
        "    denominator = noise_power + inter_user_interference\n",
        "\n",
        "    result = numerator/denominator\n",
        "    return result\n",
        "\n",
        "\n",
        "def compute_WSR_nn(user_weights, channel, precoder, noise_power, nr_of_users):\n",
        "\n",
        "   result = 0\n",
        "   for batch_index in range(nr_of_samples_per_batch):\n",
        "    for user_index in range(nr_of_users):\n",
        "        user_sinr = compute_sinr_nn(channel[batch_index], precoder[batch_index], noise_power, user_index,nr_of_users)\n",
        "        result = result + user_weights[batch_index][user_index]*(tf.log(1 + user_sinr)/tf.log(tf.cast(2.0,tf.float64)))\n",
        "\n",
        "   return result\n",
        "\n",
        "\n",
        "# Computes a channel realization and returns it in two formats, one for the WMMSE and one for the unfolded WMMSE.\n",
        "# It also returns the initialization value of the transmitter precoder, which is used as input in the computation graph of the unfolded WMMSE.\n",
        "def compute_channel(nr_of_BS_antennas, nr_of_users, total_power):\n",
        "  channel_nn = []\n",
        "  initial_transmitter_precoder = []\n",
        "  channel_WMMSE = np.zeros((nr_of_users, nr_of_BS_antennas)) + 1j*np.zeros((nr_of_users, nr_of_BS_antennas))\n",
        "\n",
        "  for i in range(nr_of_users):\n",
        "      result_real = np.sqrt(0.5)*np.random.normal(size = (nr_of_BS_antennas,1))\n",
        "      result_imag  =  np.sqrt(0.5)*np.random.normal(size = (nr_of_BS_antennas,1))\n",
        "      channel_WMMSE[i,:] = np.reshape(result_real,(1,nr_of_BS_antennas)) + 1j*np.reshape(result_imag, (1,nr_of_BS_antennas))\n",
        "      result_col_1 = np.vstack((result_real,result_imag))\n",
        "      result_col_2 = np.vstack((-result_imag,result_real))\n",
        "      result =  np.hstack((result_col_1, result_col_2))\n",
        "      initial_transmitter_precoder.append(result_col_1)\n",
        "      channel_nn.append(result)\n",
        "\n",
        "  initial_transmitter_precoder_array = np.array(initial_transmitter_precoder)\n",
        "  initial_transmitter_precoder_array = np.sqrt(total_power)*initial_transmitter_precoder_array/np.linalg.norm(initial_transmitter_precoder_array)\n",
        "  initial_transmitter_precoder = []\n",
        "\n",
        "  for i in range(nr_of_users):\n",
        "    initial_transmitter_precoder.append(initial_transmitter_precoder_array[i])\n",
        "\n",
        "  return channel_nn, initial_transmitter_precoder, channel_WMMSE\n",
        "\n",
        "\n",
        "def layer_pg( init, name, mse_weights, user_weights, receiver_precoder,channel, initial_transmitter_precoder, total_power):\n",
        "\n",
        "    with tf.variable_scope(name): \n",
        "\n",
        "      step_size =  tf.Variable(tf.constant(init, dtype=tf.float64), name=\"step_size\", dtype=tf.float64)\n",
        "\n",
        "      # First iteration\n",
        "      a1_exp = tf.tile(tf.expand_dims(mse_weights[:,0,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "      a2_exp = tf.tile(tf.expand_dims(user_weights[:,0,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "      a3_exp = tf.tile(tf.expand_dims((tf.reduce_sum(receiver_precoder[:,0,:,:],axis =-2))**2,-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])    \n",
        "      temp = a1_exp*a2_exp*a3_exp*tf.matmul(channel[:,0,:,:],tf.transpose(channel[:,0,:,:],perm = [0,2,1]))\n",
        "      \n",
        "      # Next iterations\n",
        "      for i in range(1, nr_of_users):\n",
        "        a1_exp = tf.tile(tf.expand_dims(mse_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "        a2_exp = tf.tile(tf.expand_dims(user_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "        a3_exp = tf.tile(tf.expand_dims((tf.reduce_sum(receiver_precoder[:,i,:,:],axis =-2))**2,-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "        temp = temp + a1_exp*a2_exp*a3_exp*tf.matmul(channel[:,i,:,:],tf.transpose(channel[:,i,:,:],perm = [0,2,1]))\n",
        "\n",
        "      sum_gradient = temp # it corresponds to eq. (6) in our paper\n",
        "\n",
        "      gradient = []\n",
        "\n",
        "      # Gradient computation\n",
        "      for i in range(nr_of_users):\n",
        "        a1_exp = tf.tile(tf.expand_dims(mse_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,1])\n",
        "        a2_exp = tf.tile(tf.expand_dims(user_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,1])\n",
        "        gradient.append(step_size*(-2.0*a1_exp*a2_exp*tf.matmul(channel[:,i,:,:],receiver_precoder[:,i,:,:])+ 2*tf.matmul(sum_gradient,initial_transmitter_precoder[:,i,:,:]))) \n",
        "        \n",
        "      tf.stack(gradient)\n",
        "      gradient = tf.transpose( tf.stack(gradient), perm=[1,0,2,3])\n",
        "      output_temp = initial_transmitter_precoder - gradient # it corresponds to eq. (12a) in our paper\n",
        "\n",
        "      output = []\n",
        "      for i in range(nr_of_samples_per_batch):\n",
        "        output.append(tf.cond((tf.linalg.norm(output_temp[i]))**2 < total_power, lambda: output_temp[i] , lambda: tf.sqrt(tf.cast(total_power, tf.float64))*output_temp[i]/tf.linalg.norm(output_temp[i]))) # it corresponds to eq. (13) in our paper\n",
        "\n",
        "      return tf.stack(output), step_size\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q37r-dBd-0K1",
        "colab_type": "text"
      },
      "source": [
        "### Function to run the WMMSE algorithm described in the paper by Shi et al."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F8me_JRP-0bw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def run_WMMSE(epsilon, channel, selected_users, total_power, noise_power, user_weights, max_nr_of_iterations, log = False):\n",
        "\n",
        "  nr_of_users = np.size(channel,0)\n",
        "  nr_of_BS_antennas = np.size(channel,1)\n",
        "  WSR=[] # to check if the WSR (our cost function) increases at each iteration of the WMMSE\n",
        "  break_condition = epsilon + 1 # break condition to stop the WMMSE iterations and exit the while\n",
        "  receiver_precoder = np.zeros(nr_of_users) + 1j*np.zeros(nr_of_users) # receiver_precoder is \"u\" in the paper of Shi et al. (it's a an array of complex scalars)\n",
        "  mse_weights = np.ones(nr_of_users) # mse_weights is \"w\" in the paper of Shi et al. (it's a an array of real scalars)\n",
        "  transmitter_precoder = np.zeros((nr_of_users, nr_of_BS_antennas)) + 1j*np.zeros((nr_of_users, nr_of_BS_antennas))# transmitter_precoder is \"v\" in the paper of Shi et al. (it's a complex matrix)\n",
        "  \n",
        "  new_receiver_precoder = np.zeros(nr_of_users) + 1j*np.zeros(nr_of_users) # for the first iteration \n",
        "  new_mse_weights = np.zeros(nr_of_users) # for the first iteration\n",
        "  new_transmitter_precoder = np.zeros((nr_of_users, nr_of_BS_antennas)) + 1j*np.zeros((nr_of_users, nr_of_BS_antennas)) # for the first iteration\n",
        "\n",
        "  \n",
        "  # Initialization of transmitter precoder\n",
        "  for user_index in range(nr_of_users):\n",
        "    if user_index in selected_users:\n",
        "      transmitter_precoder[user_index,:] = channel[user_index,:]\n",
        "  transmitter_precoder = transmitter_precoder/np.linalg.norm(transmitter_precoder)*np.sqrt(total_power)\n",
        "  \n",
        "  # Store the WSR obtained with the initialized trasmitter precoder    \n",
        "  WSR.append(compute_weighted_sum_rate(user_weights, channel, transmitter_precoder, noise_power, selected_users))\n",
        "\n",
        "  # Compute the initial power of the transmitter precoder\n",
        "  initial_power = 0\n",
        "  for user_index in range(nr_of_users):\n",
        "    if user_index in selected_users:\n",
        "      initial_power = initial_power + (compute_norm_of_complex_array(transmitter_precoder[user_index,:]))**2 \n",
        "  if log == True:\n",
        "    print(\"Power of the initialized transmitter precoder:\", initial_power)\n",
        "\n",
        "  nr_of_iteration_counter = 0 # to keep track of the number of iteration of the WMMSE\n",
        "\n",
        "  while break_condition >= epsilon and nr_of_iteration_counter<=max_nr_of_iterations:\n",
        "    \n",
        "    nr_of_iteration_counter = nr_of_iteration_counter + 1\n",
        "    if log == True:\n",
        "      print(\"WMMSE ITERATION: \", nr_of_iteration_counter)\n",
        "\n",
        "    # Optimize receiver precoder - eq. (5) in the paper of Shi et al.\n",
        "    for user_index_1 in range(nr_of_users):\n",
        "      if user_index_1 in selected_users:\n",
        "        user_interference = 0.0\n",
        "        for user_index_2 in range(nr_of_users):\n",
        "          if user_index_2 in selected_users:\n",
        "            user_interference = user_interference + (np.absolute(np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_2,:])))**2\n",
        "\n",
        "        new_receiver_precoder[user_index_1] = np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_1,:]) / (noise_power + user_interference)\n",
        "\n",
        "    # Optimize mse_weights - eq. (13) in the paper of Shi et al.\n",
        "    for user_index_1 in range(nr_of_users):\n",
        "      if user_index_1 in selected_users:\n",
        "\n",
        "        user_interference = 0 # it includes the channel of all selected users\n",
        "        inter_user_interference = 0 # it includes the channel of all selected users apart from the current one\n",
        "        \n",
        "        for user_index_2 in range(nr_of_users):\n",
        "          if user_index_2 in selected_users:\n",
        "            user_interference = user_interference + (np.absolute(np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_2,:])))**2\n",
        "        for user_index_2 in range(nr_of_users):\n",
        "          if user_index_2 != user_index_1 and user_index_2 in selected_users:\n",
        "            inter_user_interference = inter_user_interference + (np.absolute(np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_2,:])))**2\n",
        "        \n",
        "        new_mse_weights[user_index_1] = (noise_power + user_interference)/(noise_power + inter_user_interference)\n",
        "\n",
        "    A = np.zeros((nr_of_BS_antennas,nr_of_BS_antennas))+1j*np.zeros((nr_of_BS_antennas,nr_of_BS_antennas))\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:\n",
        "        # hh should be an hermitian matrix of size (nr_of_BS_antennas X nr_of_BS_antennas)\n",
        "        hh = np.matmul(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)),np.conj(np.transpose(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)))))\n",
        "        A = A + (new_mse_weights[user_index]*user_weights[user_index]*(np.absolute(new_receiver_precoder[user_index]))**2)*hh\n",
        "\n",
        "    Sigma_diag_elements_true, U = np.linalg.eigh(A)\n",
        "    Sigma_diag_elements = copy.deepcopy(np.real(Sigma_diag_elements_true))\n",
        "    Lambda = np.zeros((nr_of_BS_antennas,nr_of_BS_antennas)) + 1j*np.zeros((nr_of_BS_antennas,nr_of_BS_antennas))\n",
        "    \n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:     \n",
        "        hh = np.matmul(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)),np.conj(np.transpose(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)))))\n",
        "        Lambda = Lambda + ((user_weights[user_index])**2)*((new_mse_weights[user_index])**2)*((np.absolute(new_receiver_precoder[user_index]))**2)*hh\n",
        "\n",
        "    Phi = np.matmul(np.matmul(np.conj(np.transpose(U)),Lambda),U)\n",
        "    Phi_diag_elements_true = np.diag(Phi)\n",
        "    Phi_diag_elements = copy.deepcopy(Phi_diag_elements_true)\n",
        "    Phi_diag_elements = np.real(Phi_diag_elements)\n",
        "\n",
        "    for i in range(len(Phi_diag_elements)):\n",
        "      if Phi_diag_elements[i]<np.finfo(float).eps:\n",
        "        Phi_diag_elements[i] = np.finfo(float).eps\n",
        "      if (Sigma_diag_elements[i])<np.finfo(float).eps:\n",
        "        Sigma_diag_elements[i] = 0\n",
        "\n",
        "    # Check if mu = 0 is a solution (eq.s (15) and (16) of in the paper of Shi et al.)\n",
        "    power = 0 # the power of transmitter precoder (i.e. sum of the squared norm)\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:\n",
        "        if np.linalg.det(A) != 0:\n",
        "          w = np.matmul(np.linalg.inv(A),np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)))*user_weights[user_index]*new_mse_weights[user_index]*(new_receiver_precoder[user_index])\n",
        "          power = power + (compute_norm_of_complex_array(w))**2\n",
        "\n",
        "    # If mu = 0 is a solution, then mu_star = 0\n",
        "    if np.linalg.det(A) != 0 and power <= total_power:\n",
        "      mu_star = 0\n",
        "    # If mu = 0 is not a solution then we search for the \"optimal\" mu by bisection\n",
        "    else:\n",
        "      power_distance = [] # list to store the distance from total_power in the bisection algorithm \n",
        "      mu_low = np.sqrt(1/total_power*np.sum(Phi_diag_elements))\n",
        "      mu_high = 0\n",
        "      low_point = compute_P(Phi_diag_elements, Sigma_diag_elements, mu_low)\n",
        "      high_point = compute_P(Phi_diag_elements, Sigma_diag_elements, mu_high)\n",
        "\n",
        "      obtained_power = total_power + 2*power_tolerance # initialization of the obtained power such that we enter the while \n",
        "\n",
        "      # Bisection search\n",
        "      while np.absolute(total_power - obtained_power) > power_tolerance:\n",
        "        mu_new = (mu_high + mu_low)/2\n",
        "        obtained_power = compute_P(Phi_diag_elements, Sigma_diag_elements, mu_new) # eq. (18) in the paper of Shi et al.\n",
        "        power_distance.append(np.absolute(total_power - obtained_power))\n",
        "        if obtained_power > total_power:\n",
        "          mu_high = mu_new\n",
        "        if obtained_power < total_power:\n",
        "          mu_low = mu_new\n",
        "      mu_star = mu_new\n",
        "      if log == True:\n",
        "        print(\"first value:\", power_distance[0])\n",
        "        plt.title(\"Distance from the target value in bisection (it should decrease)\")\n",
        "        plt.plot(power_distance)\n",
        "        plt.show()\n",
        "\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:\n",
        "        new_transmitter_precoder[user_index,:] = np.matmul(np.linalg.inv(A + mu_star*np.eye(nr_of_BS_antennas)),channel[user_index,:])*user_weights[user_index]*new_mse_weights[user_index]*(new_receiver_precoder[user_index]) \n",
        "\n",
        "    # To select only the weights of the selected users to check the break condition\n",
        "    mse_weights_selected_users = []\n",
        "    new_mse_weights_selected_users = []\n",
        "    for user_index in range(nr_of_users): \n",
        "      if user_index in selected_users:\n",
        "        mse_weights_selected_users.append(mse_weights[user_index])\n",
        "        new_mse_weights_selected_users.append(new_mse_weights[user_index])\n",
        "\n",
        "    mse_weights = deepcopy(new_mse_weights)\n",
        "    transmitter_precoder = deepcopy(new_transmitter_precoder)\n",
        "    receiver_precoder = deepcopy(new_receiver_precoder)\n",
        "\n",
        "    WSR.append(compute_weighted_sum_rate(user_weights, channel, transmitter_precoder, noise_power, selected_users))\n",
        "    break_condition = np.absolute(np.log2(np.prod(new_mse_weights_selected_users))-np.log2(np.prod(mse_weights_selected_users)))\n",
        "\n",
        "  if log == True:\n",
        "    plt.title(\"Change of the WSR at each iteration of the WMMSE (it should increase)\")\n",
        "    plt.plot(WSR,'bo')\n",
        "    plt.show()\n",
        "\n",
        "  return transmitter_precoder, receiver_precoder, mse_weights, WSR[-1]\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MwsPve2s-1f_",
        "colab_type": "text"
      },
      "source": [
        "### Tensorflow computation graph to run the unfolded WMMSE algorithm described in our paper"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LSbUVrNg-1r6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "tf.reset_default_graph()\n",
        "\n",
        "channel_input = tf.placeholder(tf.float64, shape = None, name = 'channel_input')\n",
        "initial_tp = tf.placeholder(tf.float64, shape = None, name = 'initial_transmitter_precoder')\n",
        "initial_transmitter_precoder = initial_tp\n",
        "\n",
        "# Arrays that contain the initialization values of the step sizes.\n",
        "# The number of step sizes depends on the selected nuber of PGD layers, the number of elements for each step size initializer depends on the selected number of unfolded iterations\n",
        "step_size_init_1 = [1.0]\n",
        "step_size_init_2 = [1.0]\n",
        "step_size_init_3 = [1.0]\n",
        "step_size_init_4 = [1.0]\n",
        "\n",
        "# Used to collect the step sizes at each iteration\n",
        "all_step_size1_temp = []\n",
        "all_step_size2_temp = []\n",
        "all_step_size3_temp = []\n",
        "all_step_size4_temp = []\n",
        "\n",
        "profit = [] # stores the WSR obtained at each iteration\n",
        "\n",
        "for loop in range(0,nr_of_iterations_nn):\n",
        "\n",
        "  user_interference2 = []\n",
        "\n",
        "  for batch_index in range(nr_of_samples_per_batch): \n",
        "    user_interference_single = []\n",
        "    for i in range(nr_of_users):\n",
        "      temp = 0.0\n",
        "      for j in range(nr_of_users):\n",
        "        temp = temp + tf.reduce_sum((tf.matmul(tf.transpose(channel_input[batch_index, i,:,:]),initial_transmitter_precoder[batch_index,j,:,:]))**2)\n",
        "      user_interference_single.append(temp + noise_power) \n",
        "    user_interference2.append(user_interference_single)\n",
        "    \n",
        "  tf.stack(user_interference2)\n",
        "\n",
        "  user_interference_exp2 = tf.tile(tf.expand_dims(tf.tile(tf.expand_dims(user_interference2,-1),[1,1,2]),-1),[1,1,1,1])\n",
        "\n",
        "  receiver_precoder_temp = (tf.matmul(tf.transpose(channel_input,perm = [0,1,3,2]),initial_transmitter_precoder))\n",
        "  # Optimize the receiver precoder - eq. (5b) in our paper\n",
        "  receiver_precoder = tf.divide(receiver_precoder_temp,user_interference_exp2)\n",
        "\n",
        "  # Optimize the mmse weights - eq. (5a) in our paper\n",
        "  self_interference = tf.reduce_sum((tf.matmul(tf.transpose(channel_input,perm = [0,1,3,2]),initial_transmitter_precoder))**2, axis = 2)\n",
        "\n",
        "  inter_user_interference_total = []\n",
        "\n",
        "  for batch_index in range(nr_of_samples_per_batch):\n",
        "    inter_user_interference_temp = []\n",
        "    for i in range(nr_of_users):\n",
        "      temp = 0.0\n",
        "      for j in range(nr_of_users):\n",
        "        if j != i:\n",
        "          temp = temp + tf.reduce_sum((tf.matmul(tf.transpose(channel_input[batch_index, i,:,:]),initial_transmitter_precoder[batch_index,j,:,:]))**2)\n",
        "      inter_user_interference_temp.append(temp + noise_power) # $sum{|(h_i)*H,v_i}|**2 + noise_power$\n",
        "    inter_user_interference = tf.reshape(tf.stack(inter_user_interference_temp),(nr_of_users,1)) # Nx1 $sum{|(h_i)*H,v_i}|**2 + noise_power$\n",
        "    inter_user_interference_total.append(inter_user_interference)\n",
        "      \n",
        "  mse_weights = (tf.divide(self_interference,inter_user_interference_total)) + 1.0\n",
        "\n",
        "  # Optimize the transmitter precoder through PGD (here denoted as \"pg\")\n",
        "  transmitter_precoder1, step_size1 = layer_pg( step_size_init_1[loop],'pg1', mse_weights, user_weights, receiver_precoder,channel_input, initial_transmitter_precoder, total_power)\n",
        "  transmitter_precoder2, step_size2 = layer_pg( step_size_init_2[loop],'pg2', mse_weights, user_weights, receiver_precoder,channel_input, transmitter_precoder1, total_power)\n",
        "  transmitter_precoder3, step_size3 = layer_pg( step_size_init_3[loop],'pg3', mse_weights, user_weights, receiver_precoder,channel_input, transmitter_precoder2, total_power)\n",
        "  transmitter_precoder, step_size4 = layer_pg( step_size_init_4[loop],'pg4', mse_weights, user_weights, receiver_precoder,channel_input, transmitter_precoder3, total_power)\n",
        "\n",
        "  initial_transmitter_precoder = transmitter_precoder\n",
        "\n",
        "  all_step_size1_temp.append(step_size1)\n",
        "  all_step_size2_temp.append(step_size2)\n",
        "  all_step_size3_temp.append(step_size3)\n",
        "  all_step_size4_temp.append(step_size4)\n",
        "\n",
        "\n",
        "  # The WSR achieved with the transmitter precoder obtained at the current iteration is appended\n",
        "  profit.append(compute_WSR_nn(user_weights, channel_input, initial_transmitter_precoder, noise_power,nr_of_users))\n",
        "\n",
        "all_step_size1 = tf.stack(all_step_size1_temp)\n",
        "all_step_size2 = tf.stack(all_step_size2_temp)\n",
        "all_step_size3 = tf.stack(all_step_size3_temp)\n",
        "all_step_size4 = tf.stack(all_step_size4_temp)\n",
        "\n",
        "\n",
        "final_precoder = initial_transmitter_precoder # this is the last transmitter precoder, i.e. the one that will be actually used for transmission\n",
        "\n",
        "WSR = tf.reduce_sum(tf.stack(profit)) # this is the cost function to maximize, i.e. the WSR obtained if we use the transmitter precoder that we have at each round of the loop \n",
        "WSR_final = compute_WSR_nn(user_weights, channel_input, final_precoder, noise_power,nr_of_users)/nr_of_samples_per_batch # this is the WSR computed using the \"final_precoder\"\n",
        "\n",
        "optimizer = tf.train.AdamOptimizer(learning_rate=0.001).minimize(-WSR)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IbxrxLHm_FYp",
        "colab_type": "text"
      },
      "source": [
        "### Running the unfolded WMMSE and the WMMSE\n",
        "In this code we train the neural network and we compare its performance with the WMMSE, by testing both algorithms on the same dataset of channel realizations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dtZIqgqM_FkX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "WSR_WMMSE =[]\n",
        "WSR_nn = []\n",
        "training_loss = []\n",
        "\n",
        "with tf.Session() as sess:\n",
        "    print(\"start of session\")\n",
        "    start_of_time = time.time()\n",
        "    sess.run(tf.global_variables_initializer())\n",
        "    \n",
        "    for i in range(nr_of_batches_training):\n",
        "      batch_for_training = []\n",
        "      initial_transmitter_precoder_batch = []\n",
        "\n",
        "      # Building a batch for training\n",
        "      for ii in range(nr_of_samples_per_batch):\n",
        "        channel_realization_nn,init_transmitter_precoder,_ = compute_channel(nr_of_BS_antennas, nr_of_users,total_power)\n",
        "        batch_for_training.append(channel_realization_nn)\n",
        "        initial_transmitter_precoder_batch.append(init_transmitter_precoder)\n",
        "      # Training\n",
        "      sess.run(optimizer, feed_dict={channel_input:batch_for_training,initial_tp:initial_transmitter_precoder_batch })\n",
        "      training_loss.append(-1*(sess.run(WSR, feed_dict={channel_input:batch_for_training,initial_tp:initial_transmitter_precoder_batch })))\n",
        "    \n",
        "    print(\"step size1\", sess.run(all_step_size1))\n",
        "    print(\"step size2\", sess.run(all_step_size2))\n",
        "    print(\"step size3\", sess.run(all_step_size3))\n",
        "    print(\"step size4\", sess.run(all_step_size4))\n",
        "\n",
        "    print(\"Training took:\", time.time()-start_of_time)\n",
        "\n",
        "    # For repeatability\n",
        "    np.random.seed(1234)\n",
        "\n",
        "    for i in range(nr_of_batches_test):    \n",
        "      batch_for_testing = []\n",
        "      initial_transmitter_precoder_batch = []\n",
        "      WSR_WMMSE_batch = 0.0\n",
        "\n",
        "      # Building a batch for testing\n",
        "      for ii in range(nr_of_samples_per_batch):       \n",
        "        channel_realization_nn,init_transmitter_precoder,channel_realization_regular = compute_channel(nr_of_BS_antennas, nr_of_users,total_power)    \n",
        "        _,_,_,WSR_WMMSE_one_sample = run_WMMSE(epsilon, channel_realization_regular, scheduled_users, total_power, noise_power, user_weights_for_regular_WMMSE, nr_of_iterations-1, log = False)      \n",
        "        WSR_WMMSE_batch =  WSR_WMMSE_batch + WSR_WMMSE_one_sample\n",
        "        batch_for_testing.append(channel_realization_nn)\n",
        "        initial_transmitter_precoder_batch.append(init_transmitter_precoder)\n",
        "      #Testing\n",
        "      WSR_nn.append(sess.run(WSR_final , feed_dict={channel_input:batch_for_testing,initial_tp:initial_transmitter_precoder_batch }))\n",
        "      WSR_WMMSE.append(WSR_WMMSE_batch/nr_of_samples_per_batch)\n",
        "    \n",
        "print(\"Training and testing took:\", time.time()-start_of_time)\n",
        "print(\"The WSR acheived with the unfolded WMMSE algorithm is: \",np.mean(WSR_nn))\n",
        "print(\"The WSR acheived with the WMMSE algorithm is: \",np.mean(WSR_WMMSE))\n",
        "\n",
        "plt.figure()\n",
        "plt.plot(training_loss)\n",
        "plt.ylabel(\"Training loss\")\n",
        "plt.xlabel(\"Sample index\")\n"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}
